/*
  Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

  This file is part of Sunrise.

  Sunrise is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  Sunrise is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.
  
*/

// $Id$

#include<iostream>
#include "CCfits/CCfits"
#include "counter.h"
#include "polychromatic_grid.h"
#include "optical.h"
#include "scatterer.h"
#include "constants.h"
#include "dust_model.h"
#include "emission.h"
#include "emission-fits.h"
#include "full_sed_emergence.h"
#include "full_sed_grid.h"
#include "emergence-fits.h"
#include "mcrx-units.h"
#include "grid_factory.h"
#include "grid-fits.h"
#include "shoot.h"
#include "fits-utilities.h"
#include "boost/shared_ptr.hpp"
#include "density_generator.h"
#include "preferences.h"
#include "model.h"
#include "grain_model_hack.h"
#include "wd01_grain_model.h"
#include "wd01_Brent_PAH_grain_model.h"
#include "p04_grain_model.h"
#include "molecule_grain_model.h"
#include "grain_model_hack.h"
#include "xfer.h"
#include "ir_grid.h"
#include "equilibrium.h"
#include "vecops.h"
#include "blackbody_grain.h"
#include "boost/program_options.hpp"
#include "blitz-fits.h"
#include "boost/lexical_cast.hpp"
#include "imops.h"

using namespace mcrx;
using namespace CCfits;
using namespace std; 
namespace po = boost::program_options;

typedef local_random T_rng_policy;
typedef scatterer<polychromatic_scatterer_policy, T_rng_policy> T_scatterer;
typedef T_scatterer::T_lambda T_lambda;
typedef T_scatterer::T_biaser T_biaser;
typedef dust_model<T_scatterer, cumulative_sampling, T_rng_policy> T_dust_model;
typedef full_sed_emergence T_emergence;

typedef full_sed_grid<adaptive_grid> T_grid;
typedef ir_grid<adaptive_grid, cumulative_sampling, local_random> T_ir_grid;
typedef emission<polychromatic_policy, T_rng_policy> T_emission;

class uniform_factory: public grid_factory<T_grid::T_data> {
public:
  const T_densities rho_, n_tol_;
  const T_float rmax_;
  const int n_lambda, maxlevel;

  T_densities get_rho(const T_cell& c) {
    // return 0 if outside of radius rmax
    T_densities rho = (dot(c.getcenter(), c.getcenter())< rmax_*rmax_) ? rho_ : T_densities(rho_*0.);
    return rho;
  };

  uniform_factory (const T_densities& rho, T_densities n_tol, 
		   int nl, int ml, T_float rmax):
    rho_ (rho), n_lambda (nl), maxlevel(ml), n_tol_(n_tol), rmax_(rmax) {};
  
  virtual bool refine_cell_p (const T_cell& c, int level) {
    if ((level<maxlevel)&&
	(abs(rmax_*rmax_-dot(c.getcenter(),c.getcenter())) < 
	 dot(c.getsize(),c.getsize()))) // && any(get_rho(c)>0))
      return true;
      
    // cell is refined to keep column density below n_tol
    T_densities rho = get_rho(c);
    T_densities n2 (dot(c.getsize(), c.getsize())*rho*rho);
    bool refine = all( n2 > (n_tol_*n_tol_));
    //if(!refine) cout << "No refinement level " << level << " at " << c.getcenter() << c.getsize() << endl;
    return refine;
  };
  virtual bool unify_cell_p (T_grid::const_local_iterator begin,
			     T_grid::const_local_iterator end) {
    T_densities s (rho_*0.);
    T_densities ssq (rho_*0.);
    const vec3d sz=begin->getsize();

    int n=1;
    while (begin != end) {
      T_densities o = (begin++)->data()->densities();
      s+=o;
      ssq+=o*o;
      n++;
    }

    const T_densities mean (s/n);
    const T_densities stdev2 (ssq/n- mean*mean);

    // unify only if the resulting cell would be < n_tol
    T_densities n2 (dot(2*sz,2*sz)*mean*mean);
    bool unify = all( n2 < (n_tol_*n_tol_));
    
    return all (stdev2==0 ) && unify;
  };
  virtual T_data get_data (const T_cell& c) {
    return T_data(get_rho(c), vec3d(0,0,0), T_lambda(n_lambda));
  }
    
  virtual int n_threads () {return 1;};
  virtual int work_chunk_levels () {return 5;};
  virtual int estimated_levels_needed () {return 0;};
};

void ir_thick_test (po::variables_map opt)
{
  const int n_threads = opt["n_threads"].as<int>();
  const long n_rays =opt["n_rays"].as<long>();
  const long n_rays_ir =opt["n_rays_ir"].as<long>();
  const long n_rays_temp = opt["n_rays_temp"].as<long>();
  const int grid_n = 10;
  const T_float grid_extent=5;
  const int maxlevel= opt["maxlevel"].as<int>();
  const int n_forced=1;
  const int n_forced_ir=opt["n_ir_forced"].as<int>();
  const T_float vis_ref_lambda = opt["vis_ref"].as<T_float>();
  const T_float ir_ref_lambda = opt["ir_ref"].as<T_float>();
  const T_float rel_tol =opt["rel_tol"].as<T_float>();
  const T_float abs_tol =opt["abs_tol"].as<T_float>();
  const T_float tau = opt["tau"].as<T_float>();
  const T_float tau_tol = opt["tau_tol"].as<T_float>();
  const string output_file(opt["output-file"].as<string>());
  const T_float i_rr=1e-2;
  const T_float cameradist=1e3;
  const T_float source_temp = opt["source_temp"].as<T_float>();
  const T_float source_lum = opt["source_lum"].as<T_float>();

  mcrx::seed(opt["seed"].as<int>());
  T_unit_map units;
  units ["length"] = "au";
  units ["mass"] = "Msun";
  units ["luminosity"] = "W";
  units ["wavelength"] = "m";
  units ["L_lambda"] = "W/m";

  // create wavelengths
  const int n_lambdas=50;
  array_1 lambda;

  Preferences p;
  p.setValue("grain_data_directory", opt["grain_data_dir"].as<string>());
  p.setValue("wd01_parameter_set", string("DL07_MW3.1_60"));
  p.setValue("template_pah_fraction", 0.5);
  //  p.setValue("wd01_parameter_set", string("LMCavg_20"));3

  p.setValue("use_dl07_opacities", true);
  //p.setValue("use_dl07_opacities", false);
  p.setValue("n_threads", n_threads);
  p.setValue("use_grain_temp_lookup", true);
  p.setValue("use_cuda", false);

  // load grain model
  T_dust_model::T_scatterer_vector sv;
  if(opt["use_blackbody"].as<bool>()) {
    lambda.reference(logspace(2e-7,5e-3,n_lambdas));
    sv.push_back(boost::shared_ptr<T_scatterer> 
 		 (new blackbody_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>(1.0, units)));
  }
  else {
    lambda.reference(logspace(1e-7,1e-3,n_lambdas));    
    sv.push_back(boost::shared_ptr<T_scatterer> 
 		 (new wd01_Brent_PAH_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>(p, units)));
  }
  T_dust_model model (sv);
  model.set_wavelength(lambda);
  vector<T_float> lambdas(lambda.begin(), lambda.end());

  const int lambda550 = lower_bound(lambdas.begin(), lambdas.end(),
				    0.551e-6) - lambdas.begin()-1;
  // find opacity normalization for whatever dust we're using
  const T_float unit_opacity_rho = 1.0/sv[0]->opacity()(lambda550);
  T_densities rho(1); rho=unit_opacity_rho/grid_extent*tau;
  T_densities n_tol(1); n_tol = tau_tol*unit_opacity_rho;

  // build grid
  cout << "building grid" << endl;
  uniform_factory factory (rho,n_tol, lambda.size(), maxlevel, grid_extent);
  // first create the adaptive grid
  boost::shared_ptr<adaptive_grid<T_grid::T_data> > gg
    (new adaptive_grid<T_grid::T_data>  
     (vec3d (-grid_extent,-grid_extent,-grid_extent),
      vec3d (grid_extent,grid_extent,grid_extent),
      vec3i (grid_n, grid_n,grid_n), factory));

  // and then the full_sed_grid
  T_grid gr(gg, units);

  // calculate total dust mass
  T_densities mtot(1);
  mtot=0;
  for(T_grid::const_iterator c=gr.begin(); c!=gr.end(); ++c)
    mtot+=c->data()->densities()*c->volume();
  cout << "Total dust mass: " << mtot(0) << endl;
  
  cout << "setting up emission" << endl;
  blackbody bb(source_temp, 3.91e26*source_lum);

  auto_ptr<T_emission> em;
  em.reset(new pointsource_emission<polychromatic_policy, T_rng_policy> (vec3d (.01, .01, -.01), bb.emission(lambda)));
  
  cout << bb.emission(lambda); 
  cout << "setting up emergence" << endl;
  std::vector<std::pair<T_float, T_float> > cam_pos;
  cam_pos.push_back(std:: make_pair (0.1, 0.01) );
  auto_ptr<T_emergence> e(new T_emergence(cam_pos, cameradist, 2.1*grid_extent, 50));
  
  cout << "shooting" << endl;

  std::vector<T_rng::T_state> states= generate_random_states (n_threads);
  T_rng_policy rng;

  xfer<T_dust_model, T_grid, T_rng_policy> x (gr, *em, *e, model, rng, i_rr, 0,n_forced,10.);
  long n_rays_actual = 0;
  dummy_terminator t;
  const int vis_ref = lower_bound(lambdas.begin(), lambdas.end(),
				  vis_ref_lambda) - lambdas.begin()-1;
  T_biaser b(vis_ref);
  shoot (x, scatter_shooter<T_biaser> (b), t, states, n_rays,
  	 n_rays_actual, n_threads,0, "");

  T_float normalization = 1./n_rays_actual;

  {
    FITS output("!"+opt["output-file"].as<string>(), Write);
    e->write_images(output,normalization,units, "-STAR", false, false);

    // calculate integrated SEDs
    Table* iq_hdu=output.addTable("INTEGRATED_QUANTITIES",0);
    iq_hdu->addColumn(Tdouble, "lambda", 1,
		      units.get("L_lambda"));
    write(iq_hdu->column("lambda"),lambda,1);
    for (int i=0; i<cam_pos.size(); ++i) {
      integrate_image(output,
		      "CAMERA"+boost::lexical_cast<string>(i)+"-STAR",
		      4*constants::pi*cameradist*cameradist,
		      "L_lambda"+boost::lexical_cast<string>(i),
		      lambda,
		      "",
		      units,
		      dummy_terminator());
    }
  }

  // now we start the interesting stuff...
  // make dust mass vector
  array_2 dust_mass(gr.n_cells(), model.n_scatterers());
  array_2 star_intensity(gr.intensities().shape());

  int cc=0;
  const T_float area_factor = gr.area_factor();
  for(T_grid::const_iterator c=gr.begin(); c!=gr.end(); ++c, ++cc) {
    dust_mass(cc, Range::all()) = c->data()->densities()*c->volume();
    star_intensity(cc, Range::all()) =
      (gr.intensities()(cc,Range::all())*normalization/
       (4*constants::pi*c->volume()*area_factor));
  }

  // first create the adaptive grid
  boost::shared_ptr<adaptive_grid<T_ir_grid::T_data> > irgg
    (new adaptive_grid<T_ir_grid::T_data>  
     (vec3d (-grid_extent,-grid_extent,-grid_extent),
      vec3d (grid_extent,grid_extent,grid_extent),
      vec3i (grid_n, grid_n,grid_n), gr.get_structure()));

  // and then the full grid
  T_ir_grid irg(irgg, dust_mass, lambda, units);

  array_2 dust_intensity(gr.n_cells(), lambda.size());
  dust_intensity=0;

  const int ir_ref = lower_bound(lambdas.begin(), lambdas.end(),
				 ir_ref_lambda) - lambdas.begin()-1;
  T_biaser b2(ir_ref);
  scatter_shooter<T_biaser> s(b2);

  // Determine dust equilibrium
  determine_dust_equilibrium_intensities
    (model,
     gr,
     irg,
     lambda,
     star_intensity,
     dust_intensity,
     states,
     n_rays_temp,
     n_threads,
     rel_tol,
     abs_tol,
     i_rr,
     n_forced_ir,
     10.,
     s,
     t,
     false);

  // now generate output
  std::vector<grain_model<polychromatic_scatterer_policy, 
    mcrx_rng_policy>*> grain_models;
  for(int i=0; i< model.n_scatterers(); ++i) {
    grain_models.push_back
      (dynamic_cast<grain_model<polychromatic_scatterer_policy, 
       mcrx_rng_policy>*>(&model.get_scatterer(i)));
    grain_models.back()->resample(lambda);
  }

  e.reset(new T_emergence(cam_pos, cameradist, 2.1*grid_extent, 50));
  irg.calculate_SED(array_2(star_intensity + dust_intensity), grain_models,
		    t);
  gr.reset();
  xfer<T_dust_model, T_grid, T_rng_policy> x2 (gr, irg, *e, model, rng, 1e-2, 0,0,10.);

  n_rays_actual = 0;
  shoot (x2, s, t, states, n_rays_ir,
  	 n_rays_actual, n_threads,0, "");

  normalization = irg.total_emission_weight()/n_rays_actual;

  {
    FITS output(output_file, Write);
    e->write_images(output,normalization,units, "-IR", false, false);

    // calculate integrated SEDs
    for (int i=0; i<cam_pos.size(); ++i) {
      integrate_image(output,
		      "CAMERA"+boost::lexical_cast<string>(i)+"-IR",
		      4*constants::pi*cameradist*cameradist,
		      "L_lambda_ir"+boost::lexical_cast<string>(i),
		      lambda,
		      "",
		      units,
		      dummy_terminator());
    }
  }

  ofstream out("cell_seds.txt");
  for(T_ir_grid::iterator c=irg.begin(); c!=irg.end(); ++c)
    out << sqrt(dot(c->getcenter(),c->getcenter())) << '\n';
  out << "\n\n";
  for(T_ir_grid::iterator c=irg.begin(); c!=irg.end(); ++c) {
    for(int l=0; l<lambda.size(); ++l)
      out << lambda(l) << '\t' << c->data()->get_emission()(l) << '\n';
    out << "\n\n";
  }
}

int main (int argc, char** argv)
{
  // Declare the supported options.
  po::options_description desc("Optically thick cloud with central point source test problem. Allowed options");
  desc.add_options()
    ("output-file", po::value<string>(), "output file name")
    ("n_rays", po::value<long>()->default_value(1000000), "number of rays")
    ("n_rays_temp", po::value<long>()->default_value(1000000), 
     "number of rays for dust temp iteration")
    ("n_rays_ir", po::value<long>()->default_value(1000000), 
     "number of rays for final dust emission")
    ("n_threads", po::value<int>()->default_value(8), "number of threads")
    ("maxlevel", po::value<int>()->default_value(6), 
     "max grid refinement level")
    ("n_ir_forced", po::value<int>()->default_value(0), 
     "number of forced scatterings for IR emission")
    ("vis_ref", po::value<T_float>()->default_value(0.9e-6),
     "stellar radiation reference wavelength")
    ("ir_ref", po::value<T_float>()->default_value(40e-6),
     "dust emission reference wavelength")
    ("i_rr", po::value<T_float>()->default_value(1e-2),
     "intensity where Russian Roulette starts")
    ("rel_tol", po::value<T_float>()->default_value(1e-3), 
     "relative ir equilibrium tolerance")
    ("abs_tol", po::value<T_float>()->default_value(1e-3), 
     "absolute ir equilibrium tolerance")
    ("source_temp", po::value<T_float>()->default_value(5800),
     "temperature of central blackbody source")
    ("source_lum", po::value<T_float>()->default_value(1),
     "luminosity of central source (in Lsun)")
    ("radius", po::value<T_float>()->default_value(5),
     "radius of cloud (in AU)")
    ("tau", po::value<T_float>()->default_value(100),
     "V-band optical depth to center of cloud")
    ("tau_tol", po::value<T_float>()->default_value(100),
     "cell optical depth refinement tolerance")
    ("seed", po::value<int>()->default_value(42), "random number seed")
    ("use_blackbody", po::value<bool>()->default_value(false), "Use ideal blackbody grains instead of a physical dust model ")
    ("grain_data_dir", po::value<string>()->default_value("/home/patrik/dust_data/crosssections"), "grain data directory")
    ("help", "print this page")
    ;

  po::positional_options_description p;
  p.add("output-file", 1);

  po::variables_map opt;
  po::store(po::command_line_parser(argc, argv).
	    options(desc).positional(p).run(), opt);
  po::notify(opt);    

  if (opt.count("help")) {
    cout << desc << "\n";
    return 1;
  }
  if(!opt.count("output-file")) {
    cout << "Must specify output file name on command line" << endl;
    return 1;
  }
  
  cout << "Running ir_thick_test with parameters:\n";
  cout << "output-file = " << opt["output-file"].as<string>() << '\n'
       << "n_rays = " << opt["n_rays"].as<long>() << '\n'
       << "n_rays_temp = " << opt["n_rays_temp"].as<long>() << '\n'
       << "n_rays_ir = " << opt["n_rays_ir"].as<long>() << '\n'
       << "n_threads = " << opt["n_threads"].as<int>() << '\n'
       << "maxlevel = " << opt["maxlevel"].as<int>() << '\n'
       << "n_ir_forced = " << opt["n_ir_forced"].as<int>() << '\n'
       << "vis_ref = " << opt["vis_ref"].as<T_float>() << '\n'
       << "ir_ref = " << opt["ir_ref"].as<T_float>() << '\n'
       << "i_rr = " << opt["i_rr"].as<T_float>() << '\n'
       << "rel_tol = " << opt["rel_tol"].as<T_float>() << '\n'
       << "abs_tol = " << opt["abs_tol"].as<T_float>() << '\n'
       << "source_temp = " << opt["source_temp"].as<T_float>() << '\n'
       << "source_lum = " << opt["source_lum"].as<T_float>() << '\n'
       << "radius = " << opt["radius"].as<T_float>() << '\n'
       << "tau = " << opt["tau"].as<T_float>() << '\n'
       << "tau_tol = " << opt["tau_tol"].as<T_float>() << '\n'
       << "seed = " << opt["seed"].as<int>() << '\n'
       << "use_blackbody = " << opt["use_blackbody"].as<bool>() << "\n\n";

  ir_thick_test(opt);
}
